"""
数据库连接池管理器
提供高性能的数据库连接复用和管理功能
"""

import asyncio
import time
from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from contextlib import asynccontextmanager
import logging

from app.models.datasource import DatabaseType
from app.utils.encryption import decrypt_password

logger = logging.getLogger(__name__)


@dataclass
class ConnectionConfig:
    """连接配置"""
    db_type: DatabaseType
    host: str
    port: int
    database: str
    username: str
    password: str
    connection_params: Dict[str, Any]
    max_connections: int = 10
    min_connections: int = 2
    connection_timeout: int = 30
    idle_timeout: int = 300  # 5分钟空闲超时


@dataclass
class PooledConnection:
    """池化连接"""
    connection: Any
    created_at: float
    last_used_at: float
    in_use: bool = False
    
    @property
    def is_expired(self) -> bool:
        """检查连接是否过期"""
        return time.time() - self.last_used_at > 300  # 5分钟未使用则过期


class ConnectionPool:
    """数据库连接池"""
    
    def __init__(self, config: ConnectionConfig):
        """
        初始化连接池
        
        Args:
            config: 连接配置
        """
        self.config = config
        self.connections: List[PooledConnection] = []
        self.lock = asyncio.Lock()
        self._closed = False
        
        # 连接统计
        self.stats = {
            'total_created': 0,
            'total_acquired': 0,
            'total_released': 0,
            'current_active': 0,
            'current_idle': 0
        }
    
    async def _create_connection(self) -> Any:
        """创建新的数据库连接"""
        try:
            if self.config.db_type == DatabaseType.MYSQL:
                import aiomysql
                connection = await aiomysql.connect(
                    host=self.config.host,
                    port=self.config.port,
                    user=self.config.username,
                    password=self.config.password,
                    db=self.config.database,
                    connect_timeout=self.config.connection_timeout,
                    **self.config.connection_params
                )
            elif self.config.db_type == DatabaseType.POSTGRESQL:
                import asyncpg
                connection = await asyncpg.connect(
                    host=self.config.host,
                    port=self.config.port,
                    user=self.config.username,
                    password=self.config.password,
                    database=self.config.database,
                    command_timeout=self.config.connection_timeout,
                    **self.config.connection_params
                )
            else:
                raise NotImplementedError(f"数据库类型 {self.config.db_type} 暂不支持连接池")
            
            self.stats['total_created'] += 1
            logger.debug(f"创建新连接: {self.config.db_type}://{self.config.host}:{self.config.port}")
            return connection
            
        except Exception as e:
            logger.error(f"创建数据库连接失败: {e}")
            raise
    
    async def _close_connection(self, connection: Any):
        """关闭数据库连接"""
        try:
            if self.config.db_type == DatabaseType.MYSQL:
                connection.close()
            elif self.config.db_type == DatabaseType.POSTGRESQL:
                await connection.close()
            
            logger.debug(f"关闭连接: {self.config.db_type}://{self.config.host}:{self.config.port}")
            
        except Exception as e:
            logger.error(f"关闭数据库连接失败: {e}")
    
    async def _cleanup_expired_connections(self):
        """清理过期连接"""
        expired_connections = []
        
        for pooled_conn in self.connections[:]:
            if not pooled_conn.in_use and pooled_conn.is_expired:
                expired_connections.append(pooled_conn)
                self.connections.remove(pooled_conn)
        
        # 关闭过期连接
        for pooled_conn in expired_connections:
            await self._close_connection(pooled_conn.connection)
            logger.debug("清理过期连接")
    
    async def acquire(self) -> Any:
        """获取连接"""
        if self._closed:
            raise RuntimeError("连接池已关闭")
        
        async with self.lock:
            # 清理过期连接
            await self._cleanup_expired_connections()
            
            # 查找可用连接
            for pooled_conn in self.connections:
                if not pooled_conn.in_use:
                    pooled_conn.in_use = True
                    pooled_conn.last_used_at = time.time()
                    self.stats['total_acquired'] += 1
                    self.stats['current_active'] += 1
                    self.stats['current_idle'] -= 1
                    logger.debug("复用现有连接")
                    return pooled_conn.connection
            
            # 如果没有可用连接且未达到最大连接数，创建新连接
            if len(self.connections) < self.config.max_connections:
                connection = await self._create_connection()
                pooled_conn = PooledConnection(
                    connection=connection,
                    created_at=time.time(),
                    last_used_at=time.time(),
                    in_use=True
                )
                self.connections.append(pooled_conn)
                self.stats['total_acquired'] += 1
                self.stats['current_active'] += 1
                return connection
            
            # 如果达到最大连接数，等待可用连接
            raise RuntimeError(f"连接池已满，最大连接数: {self.config.max_connections}")
    
    async def release(self, connection: Any):
        """释放连接"""
        async with self.lock:
            for pooled_conn in self.connections:
                if pooled_conn.connection == connection and pooled_conn.in_use:
                    pooled_conn.in_use = False
                    pooled_conn.last_used_at = time.time()
                    self.stats['total_released'] += 1
                    self.stats['current_active'] -= 1
                    self.stats['current_idle'] += 1
                    logger.debug("释放连接")
                    return
            
            logger.warning("尝试释放未知连接")
    
    @asynccontextmanager
    async def get_connection(self):
        """上下文管理器方式获取连接"""
        connection = await self.acquire()
        try:
            yield connection
        finally:
            await self.release(connection)
    
    async def close(self):
        """关闭连接池"""
        async with self.lock:
            self._closed = True
            
            # 关闭所有连接
            for pooled_conn in self.connections:
                await self._close_connection(pooled_conn.connection)
            
            self.connections.clear()
            logger.info(f"连接池已关闭: {self.config.db_type}://{self.config.host}:{self.config.port}")
    
    def get_stats(self) -> Dict[str, Any]:
        """获取连接池统计信息"""
        return {
            **self.stats,
            'total_connections': len(self.connections),
            'config': {
                'max_connections': self.config.max_connections,
                'min_connections': self.config.min_connections,
                'connection_timeout': self.config.connection_timeout,
                'idle_timeout': self.config.idle_timeout
            }
        }


class ConnectionPoolManager:
    """连接池管理器"""
    
    def __init__(self):
        """初始化连接池管理器"""
        self.pools: Dict[str, ConnectionPool] = {}
        self.lock = asyncio.Lock()
    
    def _get_pool_key(self, config: ConnectionConfig) -> str:
        """生成连接池键"""
        return f"{config.db_type}://{config.username}@{config.host}:{config.port}/{config.database}"
    
    async def get_pool(self, config: ConnectionConfig) -> ConnectionPool:
        """获取或创建连接池"""
        pool_key = self._get_pool_key(config)
        
        async with self.lock:
            if pool_key not in self.pools:
                self.pools[pool_key] = ConnectionPool(config)
                logger.info(f"创建新连接池: {pool_key}")
            
            return self.pools[pool_key]
    
    async def close_pool(self, config: ConnectionConfig):
        """关闭指定连接池"""
        pool_key = self._get_pool_key(config)
        
        async with self.lock:
            if pool_key in self.pools:
                await self.pools[pool_key].close()
                del self.pools[pool_key]
                logger.info(f"关闭连接池: {pool_key}")
    
    async def close_all(self):
        """关闭所有连接池"""
        async with self.lock:
            for pool in self.pools.values():
                await pool.close()
            
            self.pools.clear()
            logger.info("关闭所有连接池")
    
    def get_all_stats(self) -> Dict[str, Any]:
        """获取所有连接池统计信息"""
        return {
            pool_key: pool.get_stats()
            for pool_key, pool in self.pools.items()
        }


# 全局连接池管理器实例
connection_pool_manager = ConnectionPoolManager()


async def get_pooled_connection(datasource_config: Dict[str, Any]):
    """获取池化连接的便捷函数"""
    config = ConnectionConfig(
        db_type=datasource_config['type'],
        host=datasource_config['host'],
        port=datasource_config['port'],
        database=datasource_config['database'],
        username=datasource_config['username'],
        password=datasource_config['password'],
        connection_params=datasource_config.get('connection_params', {}),
        max_connections=datasource_config.get('max_connections', 10)
    )
    
    pool = await connection_pool_manager.get_pool(config)
    return pool.get_connection()


async def cleanup_connection_pools():
    """清理连接池（应用关闭时调用）"""
    await connection_pool_manager.close_all()
